# -*- encoding: utf-8 -*-

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

mpl.rcParams['font.sans-serif'] = 'SimHei'
mpl.rcParams['axes.unicode_minus'] = False


# 构建二维函数
def f2(x, y):
    return 0.6 * (x + y) ** 2 - x * y


def hx2(x, y):
    return 0.6 * 2 * (x + y) - y


def hy2(x, y):
    return 0.6 * 2 * (x + y) - x


# 使用梯度下降法进行解答
GD_X1 = []
GD_X2 = []
GD_Y = []
x1 = 4
x2 = 4
alpha = 0.5
f_change = f2(x1, x2)
f_current = f_change
GD_X1.append(x1)
GD_X2.append(x2)
GD_Y.append(f_current)

# 迭代次数
iter_num = 0
while f_change > 1e-10 and iter_num < 100:
    iter_num += 1
    prex1 = x1
    prex2 = x2
    x1 = x1 - alpha * hx2(prex1, prex2)
    x2 = x2 - alpha * hy2(prex1, prex2)
    tmp = f2(x1, x2)
    # 判断y的变化不能太小，不然意义不大
    f_change = np.abs(tmp - f_current)
    f_current = tmp
    GD_X1.append(x1)
    GD_X2.append(x2)
    GD_Y.append(f_current)
print(u'最终的结果：(%.5f, %.5f， %.5f)' % (x1, x2, f_current))
print(u'迭代次数是：%d' % iter_num)
print(GD_X1)
print(GD_X2)

# 构建画图数据
X1 = np.arange(-4, 4.5, 0.2)
X2 = np.arange(-4, 4.5, 0.2)

X1, X2 = np.meshgrid(X1, X2)
Y = np.array(list(map(lambda t: f2(t[0], t[1]), zip(X1.flatten(), X2.flatten()))))
Y.shape = X1.shape

# 画图
flg = plt.figure(facecolor='w')
ax = Axes3D(flg)
ax.plot_surface(X1, X2, Y, rstride=1, cstride=1, cmap=plt.cm.jet)
ax.plot(GD_X1, GD_X2, GD_Y, 'bo--', linewidth=2)

ax.set_title(u'函数$y=0.6 * (x1 + x2)^2 - x1 * x2$;\n学习率： %.3f; 最终解：(%.3f, %.3f, %.3f); 迭代次数：%d' % (
    alpha, x1, x2, f_current, iter_num))
plt.show()
